# -*- coding: utf-8 -*-
import pymongo
import urllib.parse

from FDC_spider.settings import MONGODB

"""
MONGODB = {
    'host': '192.168.1.159',
    'port': '27017',
    'user': '',
    'pwd': '',
    'replicaSet': {
        'name': '',
        "members": [
            {
                "host": "localhost",
                "port": "27017"
            },
        ]
    }
}
"""


class MongodbClient(object):
    """
    MongoDB 封装类
    """

    def __init__(self, db_name=None, col_name=None):
        self._client = self.__get_client()
        if db_name:
            self._db = self._client[db_name]
            if col_name:
                self._collection = self._client[db_name][col_name]

    def __get_client(self):
        """
        创建连接对象
        :return:
        """
        if MONGODB['replicaSet']['name']:
            host_opt = []
            for m in MONGODB['replicaSet']['members']:
                host_opt.append('%s:%s' % (m['host'], m['port']))
            replicaSet = MONGODB['replicaSet']['name']
        else:
            host_opt = '%s:%s' % (MONGODB['host'], MONGODB['port'])
            replicaSet = None

        option = {
            'host': host_opt,
            'replicaSet': replicaSet,
        }
        if MONGODB['user'] and MONGODB['pwd']:
            option['username'] = urllib.parse.quote_plus(MONGODB['user'])
            option['password'] = urllib.parse.quote_plus(MONGODB['pwd'])
            option['authMechanism'] = 'SCRAM-SHA-1'

        return pymongo.MongoClient(**option)

    def set_db(self, db_name):
        """
        设置数据库名
        :param db_name:
        :return:
        """
        self._db = self._client[db_name]

    def set_col(self, db_col):
        """
        设置集合名
        :param db_col:
        :return:
        """
        if self._db:
            self._collection = self._db[db_col]

    def insert_one(self, document):
        """
        添加一条数据
        :param document:
        :return:
        """
        ret = self._collection.insert_one(document)
        # print(ret)

    def insert_many(self, docs):
        """
        添加多条数据
        :param docs:  [{},{},{}....]
        :return:
        """
        ret = self._collection.insert_many(docs)
        return ret

    def save(self, doc, manipulate=True, check_keys=True, **kwargs):
        """
        添加/保存 一条数据
        :param doc:
        :param manipulate:    True-重复后覆盖    False-重复后新增
        :param check_keys:
        :param kwargs:
        :return:
        """
        ret = self._collection.save(doc, manipulate, check_keys, **kwargs)
        return ret

    def update_one(self, filter, document):
        """
        更新一条数据(存在就更新，不存在就插入)
        :param filter:
        :param document:
        :return:
        """
        ret = self._collection.update_one(filter, {'$set': document}, upsert=True)
        print(ret)

    def close(self):
        """
        关闭连接
        :return:
        """
        self._client.close()

    def aggregate(self, pipeline):
        """
        聚合操作
        :param pipeline: [
                            {$sort:{_id:1}},
                            {$group: { _id: {name: '$name',age: '$age'},count: {$sum: 1},dups: {$addToSet: '$_id'}}},
                            {$match: {count: {$gt: 0}}}
                        ]
        :return:
        """
        return self._collection.aggregate(pipeline)

    def delete_one(self, filter):
        """
        删除一条数据
        :param filter:
        :return:
        """
        ret = self._collection.delete_one(filter)
        return ret

    def delete_many(self, filter):
        """
        删除多条数据
        :param filter:
        :return:
        """
        ret = self._collection.delete_many(filter)
        return ret
